package com.jingding.procurement.annotation;

import com.jingding.procurement.ex.ServiceException;
import com.jingding.procurement.restful.ServiceCode;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;

@Aspect
@Component
public  class RateLimitAspect {

    private Map<String,  Integer> requestCounts  =  new ConcurrentHashMap<>();

    // 环绕通知方法，在方法上加上@rateLimit注解时会触发
    @Around("@annotation(rateLimit)")
    public Object rateLimit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        // 获取HttpServletRequest对象
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest();
        // 获取请求的IP地址
        String ipAddress = request.getRemoteAddr();
        // 如果请求计数包含该IP地址并且请求次数大于等于最大限制次数
        if (requestCounts.containsKey(ipAddress) && requestCounts.get(ipAddress) >= rateLimit.maxRequests()) {
            // 抛出自定义异常，表示IP地址访问受限
            throw new ServiceException(ServiceCode.IP_LIMIT.getCode(), ServiceCode.IP_LIMIT.getMessage());
        }
        // 将请求计数的IP地址对应的请求次数+1
        requestCounts.put(ipAddress, requestCounts.getOrDefault(ipAddress, 0) + 1);
        // 创建一个定时器，根据设定的时间周期移除请求计数
        new Timer().schedule(
                new TimerTask() {
                    @Override
                    public void run() {
                        requestCounts.remove(ipAddress);
                    }
                },
                rateLimit.timePeriod()
        );
        // 执行目标方法
        return joinPoint.proceed();
    }
}
